# -*- coding: utf-8 -*-
"""
******* 文档说明 ******
基于 tkinter 写的 GUI 工具
可快速查看图片、打标 并导出配置文件

空格键：   开始、暂停播放图片
Ctrl + R： 自动播放图片顺序逆转
Ctrl + S： 保存标注标签 JSON 文件
1、2、3 ： 根据 label_mapping 打标
方向键盘： 向上、向下查看图片

# 当前项目: ImageView
# 创建时间: 2019/6/30 21:25
# 开发作者: vincent
# 创建平台: PyCharm Community Edition
# 版    本: V1.0
"""
import os
import json
import time
from tkinter import *
from PIL import Image, ImageTk


# 图片查看、标注
class AppUI(object):
    def __init__(self, img_list, label_mapping=None, scaling=0.8, sleep_time=0.5, json_path=None):
        """
        :param img_list:          图片路径列表
        :param label_mapping:     标签类别字典  '1': 'True', '2': 'False'}
        :param scaling:           图片显示缩放比例
        :param sleep_time:        图片自动播放时，图片间暂停时间
        :param json_path:         图片结果存储  JSON 文件
        """
        # 待分析图片列表（路径）
        self.img_list = img_list
        # 当前分析图片位置索引
        self.img_index = 0
        # 图片数量
        self.img_len = len(self.img_list)
        # 图片显示缩放比例
        self.scaling = scaling
        # 标签值
        self.label_mapping = label_mapping
        self.label_mapping_string = ' '.join(['{}: {}     '.format(key_i, self.label_mapping[key_i])
                                              for key_i in sorted(self.label_mapping)])
        # 标签值索引
        self.label_mapping_index = {value_: key_ for key_, value_ in self.label_mapping.items()}

        # 颜色列表
        self.color_list = ['#FFB6C1', '#FFC0CB', '#DC143C', '#FFF0F5', '#DB7093', '#FF69B4',
                           '#FF1493', '#C71585', '#DA70D6', '#D8BFD8', '#DDA0DD', '#EE82EE',
                           '#FF00FF', '#FF00FF', '#8B008B', '#800080', '#BA55D3', '#9400D3',
                           '#9932CC', '#4B0082', '#8A2BE2', '#9370DB', '#7B68EE', '#6A5ACD',
                           '#483D8B', '#E6E6FA', '#F8F8FF', '#0000FF', '#0000CD', '#191970',
                           '#00008B', '#000080', '#4169E1', '#6495ED', '#B0C4DE', '#778899',
                           '#708090', '#1E90FF', '#F0F8FF', '#4682B4', '#87CEFA', '#87CEEB',
                           '#00BFFF', '#ADD8E6', '#B0E0E6', '#5F9EA0', '#F0FFFF', '#E0FFFF',
                           '#AFEEEE', '#00FFFF', '#00FFFF', '#00CED1', '#2F4F4F', '#008B8B',
                           '#008080', '#48D1CC', '#20B2AA', '#40E0D0', '#7FFFD4', '#66CDAA',
                           '#00FA9A', '#F5FFFA', '#00FF7F', '#3CB371', '#2E8B57', '#F0FFF0',
                           '#90EE90', '#98FB98', '#8FBC8F', '#32CD32', '#00FF00', '#228B22',
                           '#008000', '#006400', '#7FFF00', '#7CFC00', '#ADFF2F', '#556B2F',
                           '#9ACD32', '#6B8E23', '#F5F5DC', '#FAFAD2', '#FFFFF0', '#FFFFE0',
                           '#FFFF00', '#808000', '#BDB76B', '#FFFACD', '#EEE8AA', '#F0E68C',
                           '#FFD700', '#FFF8DC', '#DAA520', '#B8860B', '#FFFAF0', '#FDF5E6',
                           '#F5DEB3', '#FFE4B5', '#FFA500', '#FFEFD5', '#FFEBCD', '#FFDEAD',
                           '#FAEBD7', '#D2B48C', '#DEB887', '#FFE4C4', '#FF8C00', '#FAF0E6',
                           '#CD853F', '#FFDAB9', '#F4A460', '#D2691E', '#8B4513', '#FFF5EE',
                           '#A0522D', '#FFA07A', '#FF7F50', '#FF4500', '#E9967A', '#FF6347',
                           '#FFE4E1', '#FA8072', '#FFFAFA', '#F08080', '#BC8F8F', '#CD5C5C',
                           '#FF0000', '#A52A2A', '#B22222', '#8B0000', '#800000', '#FFFFFF',
                           '#F5F5F5', '#DCDCDC', '#D3D3D3', '#C0C0C0', '#A9A9A9', '#808080']
        # 标签颜色
        label_name = [label_i[0] for label_i in sorted(self.label_mapping_index.items(), key=lambda x: x[1])]
        color_value = self.color_list[::int(len(self.color_list) // len(label_mapping))][:len(label_mapping)]
        self.label_color = dict(zip(label_name, color_value))

        # 读取标签结果，若无输入则重新创建生成
        self.json_path = json_path
        if self.json_path is None:
            self.label_dict = dict(zip(img_list, [''] * self.img_len))
        else:
            self.label_dict = json.load(open(self.json_path, "r", encoding='utf-8'))

        # 自动播放
        self.stop = True
        self.forward = True
        self.sleep_time = sleep_time

        self.root = Tk()

        # root.resizable(False, False)  # 调用方法会禁止根窗体改变大小
        # 以下方法用来计算并设置窗体显示时，在屏幕中心居中
        cur_width = self.root.winfo_width()  # get current width
        cur_height = self.root.winfo_height()  # get current height
        self.scn_width, self.scn_height = self.root.maxsize()  # get screen width and height
        window_size = '+%d+%d' % ((self.scn_width - cur_width) / 2, (self.scn_height - cur_height) / 2)
        self.root.geometry(window_size)  # 设定窗口大小

        # ########################################
        # 图片显示
        self.label_img = Label(self.root, width=int(self.scn_width * self.scaling),
                               height=int(self.scn_height * self.scaling))
        self.label_img.pack(fill=X, expand=YES, side=TOP, padx=0, pady=0)

        # ########################################
        # 标签显示
        self.btn_label = Button(self.root, text="【{:^30s}】    {}".format('', self.label_mapping_string),
                                bg='white', fg='white', width=66, height=1, state=DISABLED)
        self.btn_label.bind_all('<KeyPress>', self._keyboard_cmd)

        # 播放顺序逆转
        self.btn_label.bind_all('<Control-r>', self._auto_reverse)
        # 保存结果
        self.btn_label.bind_all('<Control-s>', self._label_save)
        self.btn_label.pack(fill=BOTH, expand=YES, padx=0, pady=0)

        # ########################################
        # 路径显示
        self.label_path = StringVar()  # 将label标签的内容设置为字符类型，用var来接收hit_me函数的传出内容用以显示在标签上
        label_path = Entry(self.root, text=self.label_path, state='readonly', bg='green', font=('Arial', 10), width=80)
        label_path.pack(fill=X, expand=YES, side=LEFT)

        # ####################################
        # 查看下一张 (快捷键 右箭头、下箭头)
        self.btn_backward = Button(self.root, text=" 》", bg='blue', fg='white', width=1, height=2,
                                   command=self.img_show_backward)
        self.btn_backward.bind_all('<KeyPress-Right>', self.img_show_backward)
        self.btn_backward.bind_all('<KeyPress-Down>', self.img_show_backward)
        self.btn_backward.pack(fill=X, expand=YES, side=LEFT)

        # 查看上一张 (快捷键 左箭头、上箭头)
        self.btn_forward = Button(self.root, text=" 《 ", bg='blue', fg='white', width=1, height=2,
                                  command=self.img_show_forward, state=DISABLED)
        self.btn_backward.bind_all('<KeyPress-Left>', self.img_show_forward)
        self.btn_backward.bind_all('<KeyPress-Up>', self.img_show_forward)
        self.btn_forward.pack(fill=X, expand=YES, side=LEFT)

        # 创建菜单栏
        # self.create_menu()
        # self.root.title("ImageView")

        # 显示第一张图片
        self._img_show()

        self.root.mainloop()

    # 创建菜单栏
    def create_menu(self):
        menu = Menu(self.root)

        about_menu = Menu(menu, tearoff=0)
        about_menu.add_command(label="图片查看打标工具 version:1.0 开发者：陈庆伟")

        # 在菜单栏中添加以下一级菜单
        menu.add_cascade(label="关于", menu=about_menu)
        self.root['menu'] = menu
        
    # 自动查看图片方向转换
    def _auto_reverse(self, *event):
        # 输入按键
        keyboard = event[0].keysym
        del keyboard

        self.forward = not self.forward
        if self.forward:
            print('AutoRun Forward!')
        else:
            print('AutoRun Backward!')

    # 根据窗口大小计算图片 Resize大小
    def _resize(self, w, h):
        # 保持原图纵横比 Resize
        if w/h > self.scn_width/self.scn_height:
            resize_w = int(self.scn_width * self.scaling)
            resize_h = int(resize_w * h / w)
        else:
            resize_h = int(self.scn_height * self.scaling)
            resize_w = int(resize_h * w / h)
        return resize_w, resize_h

    # 显示图片
    def _img_show(self):
        img_path = self.img_list[self.img_index]
        # 显示图片路径
        self.label_path.set(img_path)
        # 打开图片
        img_open = Image.open(img_path)
        # 原始大小
        raw_size = img_open.size
        # Resize
        img_open = img_open.resize(self._resize(*raw_size))
        img = ImageTk.PhotoImage(img_open)
        self.label_img.config(image=img)
        self.label_img.image = img

        # 显示图片名称及序号
        self.root.wm_title("{}/{}  {}      {}".format(
            self.img_index+1, self.img_len, os.path.basename(img_path), raw_size))

        # 显示打标标签
        text = "【{:^30s}】    {}".format(self.label_dict[self.img_list[self.img_index]], self.label_mapping_string)
        label = self.label_dict[self.img_list[self.img_index]]
        if label == '':
            bg = 'white'
        else:
            bg = self.label_color[label]
        self.btn_label.config(text=text, bg=bg)

    # 显示下张图片
    def img_show_backward(self, *event):
        if len(event) > 0:
            del event[0].keysym

        if self.img_index < self.img_len - 1:
            # 图片 索引ID +1
            self.img_index += 1
            # 显示图片
            self._img_show()

        # 无图片时按键不可用
        if self.img_index >= len(self.img_list) - 1:
            self.btn_backward.config(state=DISABLED)
        # 激活向前查看按键
        self.btn_forward.config(state=ACTIVE)

    # 显示上张图片
    def img_show_forward(self, *event):
        if len(event) > 0:
            del event[0].keysym

        if self.img_index > 0:
            # 图片 索引ID -1
            self.img_index -= 1
            # 显示图片
            self._img_show()

        # 无图片时按键不可用
        if self.img_index == 0:
            self.btn_forward.config(state=DISABLED)
        # 激活向后查看按键
        self.btn_backward.config(state=ACTIVE)

    # 键盘输入处理  打标、命令传入
    def _keyboard_cmd(self, *event):
        # 输入按键
        keyboard = event[0].keysym

        if keyboard in self.label_mapping:
            print('Label:{}'.format(self.label_mapping[keyboard]))
            # 加入当前图片标注信息
            self.label_dict[self.img_list[self.img_index]] = self.label_mapping[keyboard]

            # 图片 索引ID +1
            self.img_index += 1
            # 显示下张图片
            self._img_show()

        if keyboard == 'space':
            # 暂停开关
            self.stop = not self.stop

            if self.forward:
                for i in range(self.img_len - self.img_index - 1):
                    # 图片 索引ID +1
                    self.img_index += 1
                    # 显示下张图片
                    self._img_show()

                    self.root.update()
                    time.sleep(self.sleep_time)

                    if self.stop:
                        break
            else:
                for i in range(self.img_index):
                    # 图片 索引ID +1
                    self.img_index -= 1
                    # 显示下张图片
                    self._img_show()

                    self.root.update()
                    time.sleep(self.sleep_time)

                    if self.stop:
                        break

    # 保存标注结果
    def _label_save(self, *event):
        # 输入按键
        keyboard = event[0].keysym
        del keyboard

        time_str = time.strftime('%Y%m%d-%H%M%S', time.localtime())

        # 若无指定配置文件，则存放在当前路径下
        if self.json_path is None:
            config_name = 'Image-{}_Label-{}_{}.json'.format(self.img_len, len(self.label_mapping), time_str)
            json_path_new = os.path.join(os.path.dirname(__file__), config_name)
        # 指定配置文件，则更换当前时候保存原路径
        else:
            json_path_new = '{}_{}.json'.format('_'.join(self.json_path.split('_')[:-1]), time_str)

        with open(json_path_new, "w", encoding='utf-8') as f_json:
            json.dump(self.label_dict, f_json, indent=4, sort_keys=True, ensure_ascii=False)

        print('Label Config Save Path: {}'.format(os.path.abspath(json_path_new)))


# 配置文件图片筛选
def img_select(json_path, label, save_path=None):
    """
    :param json_path:  导出的 JSON 配置文件
    :param label:      筛选的标签类别 list
    :param save_path:  筛选图片保存路径
    :return:
    """
    import shutil

    label_dict = json.load(open(json_path, "r", encoding='utf-8'))

    if save_path is None:
        save_path = os.path.join(os.path.dirname(__file__), 'SelectImg_{}'.format(os.path.basename(json_path)[:-5]))

    os.makedirs(save_path, exist_ok=True)

    for img_path_i, label_i in label_dict.items():
        if label_i in label:
            shutil.copy(img_path_i, os.path.join(save_path, os.path.basename(img_path_i)))


if __name__ == '__main__':
    input_path = r'D:\Desktop\ImageView\image'

    # 图片路径列表
    img_path_list = [os.path.join(input_path, img_i) for img_i in os.listdir(input_path)
                     if img_i.endswith(('.jpg', '.png'))]

    # 已输出标签
    # config_json_path = r'D:\Desktop\ImageView\Image-200_Label-3_20190702-020232.json'
    config_json_path = None

    AppUI(img_path_list,   # 图片路径列表
          label_mapping={'1': 'True', '2': 'False', '3': 'Unknown'},  # 标签字典
          scaling=0.6,
          sleep_time=0,
          json_path=config_json_path)

    # # # 筛选图片另存
    # img_select(r'D:\Desktop\ImageView\Image-200_Label-3_20190702-020232.json', ['False', 'Unknown'], save_path=None)
